import torch
import os
import torch.nn as nn
import torch.nn.functional as F

import torchvision.transforms as transforms
from PIL import Image
import torchvision.models as models
# from sklearn.preprocessing import LabelBinarizer
import numpy as np

inet_pretrain = True
img_path = r"C:\Users\31956\Desktop\1.png"

nb_cls = 2

new_net = models.resnet18(pretrained=inet_pretrain)
new_net.fc = nn.Linear(512, nb_cls)

new_net.load_state_dict(torch.load(r'C:\Users\31956\Documents\coding\js\ai4-eye\eye_images\net_cataract.pt'))

new_net.eval()

if nb_cls == 2:
    classes = {0: "normal", 1: "cataract"}
elif nb_cls == 3:
    classes = {0: "normal", 1: "cataract"}
elif nb_cls == 4:
    classes = {0: "normal", 1: ""}

test_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

img = Image.open(img_path)
img_tensor = test_transform(img)

with torch.no_grad():
    outputs = new_net(img_tensor.unsqueeze(0))
    probs = F.softmax(outputs, dim=1)
    preds = torch.argmax(probs, dim=1)

print(outputs)
